import pygame
import sys
import random
import math
import time

# 初始化pygame
pygame.init()

# 游戏常量
WIDTH, HEIGHT = 800, 600
GRID_SIZE = 20
GRID_WIDTH = WIDTH // GRID_SIZE
GRID_HEIGHT = HEIGHT // GRID_SIZE
FPS = 10

# 颜色定义
BACKGROUND = (15, 20, 25)
GRID_COLOR = (30, 35, 40)
SNAKE_HEAD_COLOR = (0, 230, 118)
SNAKE_BODY_COLOR = (0, 200, 83)
ENEMY_HEAD_COLOR = (255, 50, 50)
ENEMY_BODY_COLOR = (200, 30, 30)
FOOD_COLOR = (255, 215, 0)
TEXT_COLOR = (200, 230, 255)
UI_BG = (25, 35, 45, 200)
UI_BORDER = (0, 150, 136)
SPEED_BOOST_COLOR = (100, 255, 255, 100)

# 方向常量
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

# 创建游戏窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇游戏 - 带AI敌人")
clock = pygame.time.Clock()

# 字体
font_large = pygame.font.SysFont("microsoftyahei", 48, bold=True)
font_medium = pygame.font.SysFont("microsoftyahei", 32)
font_small = pygame.font.SysFont("microsoftyahei", 24)

class Snake:
    def __init__(self, is_player=True):
        self.is_player = is_player
        self.reset()
        
    def reset(self):
        self.length = 3
        self.positions = [(random.randint(5, GRID_WIDTH-5), 
                          random.randint(5, GRID_HEIGHT-5))]
        self.direction = random.choice([UP, DOWN, LEFT, RIGHT])
        self.score = 0
        self.grow_to = 3
        self.is_alive = True
        self.last_move_time = time.time()
        # 玩家蛇初始移动速度较慢（0.2秒/格），AI蛇速度中等（0.15秒/格）
        self.move_delay = 0.2 if self.is_player else 0.15
        self.base_move_delay = self.move_delay
        
    def get_head_position(self):
        return self.positions[0]
    
    def update(self, food_positions, enemy_snake=None):
        if not self.is_alive:
            return
            
        current_time = time.time()
        if current_time - self.last_move_time < self.move_delay:
            return
            
        self.last_move_time = current_time
        
        head = self.get_head_position()
        x, y = self.direction
        new_x = (head[0] + x) % GRID_WIDTH
        new_y = (head[1] + y) % GRID_HEIGHT
        new_position = (new_x, new_y)
        
        # 检查是否撞到自己
        if new_position in self.positions[1:]:
            self.is_alive = False
            return
            
        # 检查是否撞到敌人（如果敌人存在）
        if enemy_snake and new_position in enemy_snake.positions:
            self.is_alive = False
            return
            
        self.positions.insert(0, new_position)
        
        # 如果蛇需要生长
        if len(self.positions) > self.grow_to:
            self.positions.pop()
    
    def draw(self, surface):
        for i, pos in enumerate(self.positions):
            rect = pygame.Rect(pos[0] * GRID_SIZE, pos[1] * GRID_SIZE, GRID_SIZE, GRID_SIZE)
            
            if i == 0:  # 蛇头
                color = SNAKE_HEAD_COLOR if self.is_player else ENEMY_HEAD_COLOR
                pygame.draw.rect(surface, color, rect)
                pygame.draw.rect(surface, (255, 255, 255), rect, 1)
                
                # 画眼睛
                eye_size = GRID_SIZE // 5
                eye_offset = GRID_SIZE // 4
                
                # 根据方向确定眼睛位置
                if self.direction == UP:
                    left_eye = (pos[0] * GRID_SIZE + eye_offset, pos[1] * GRID_SIZE + eye_offset)
                    right_eye = (pos[0] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size, pos[1] * GRID_SIZE + eye_offset)
                elif self.direction == DOWN:
                    left_eye = (pos[0] * GRID_SIZE + eye_offset, pos[1] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size)
                    right_eye = (pos[0] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size, pos[1] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size)
                elif self.direction == LEFT:
                    left_eye = (pos[0] * GRID_SIZE + eye_offset, pos[1] * GRID_SIZE + eye_offset)
                    right_eye = (pos[0] * GRID_SIZE + eye_offset, pos[1] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size)
                else:  # RIGHT
                    left_eye = (pos[0] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size, pos[1] * GRID_SIZE + eye_offset)
                    right_eye = (pos[0] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size, pos[1] * GRID_SIZE + GRID_SIZE - eye_offset - eye_size)
                
                pygame.draw.rect(surface, (0, 0, 0), pygame.Rect(left_eye[0], left_eye[1], eye_size, eye_size))
                pygame.draw.rect(surface, (0, 0, 0), pygame.Rect(right_eye[0], right_eye[1], eye_size, eye_size))
            else:  # 蛇身
                color = SNAKE_BODY_COLOR if self.is_player else ENEMY_BODY_COLOR
                pygame.draw.rect(surface, color, rect)
                pygame.draw.rect(surface, (0, 150, 100), rect, 1)
                
                # 添加身体纹理
                if i % 2 == 0:
                    highlight = (0, 180, 100) if self.is_player else (180, 40, 40)
                    pygame.draw.circle(surface, highlight, rect.center, GRID_SIZE // 4)
                    
    def boost_speed(self):
        # 加速效果（按住空格时）
        self.move_delay = self.base_move_delay * 0.4  # 加速到40%的原始速度
        
    def reset_speed(self):
        # 恢复原始速度
        self.move_delay = self.base_move_delay

class Food:
    def __init__(self):
        self.positions = []
        self.max_food = 5
        self.spawn_food()
        
    def spawn_food(self):
        while len(self.positions) < self.max_food:
            new_pos = (random.randint(0, GRID_WIDTH - 1), 
                       random.randint(0, GRID_HEIGHT - 1))
            self.positions.append(new_pos)
    
    def remove_food(self, position):
        if position in self.positions:
            self.positions.remove(position)
            self.spawn_food()
    
    def draw(self, surface):
        for pos in self.positions:
            rect = pygame.Rect(pos[0] * GRID_SIZE, 
                              pos[1] * GRID_SIZE, 
                              GRID_SIZE, GRID_SIZE)
            
            # 绘制食物（苹果）
            pygame.draw.circle(surface, FOOD_COLOR, rect.center, GRID_SIZE // 2)
            
            # 绘制高光
            pygame.draw.circle(surface, (255, 240, 150), 
                             (rect.centerx - 4, rect.centery - 4), 
                             GRID_SIZE // 6)

class EnemyAI:
    def __init__(self, snake, player_snake, food):
        self.snake = snake
        self.player_snake = player_snake
        self.food = food
        
    def update(self):
        if not self.snake.is_alive:
            return
            
        head = self.snake.get_head_position()
        
        # 寻找最近的食物
        nearest_food = None
        min_dist = float('inf')
        for food_pos in self.food.positions:
            dist = abs(head[0] - food_pos[0]) + abs(head[1] - food_pos[1])
            if dist < min_dist:
                min_dist = dist
                nearest_food = food_pos
                
        # 如果有食物，尝试朝食物移动
        if nearest_food:
            possible_directions = []
            
            # 计算水平和垂直方向
            dx = nearest_food[0] - head[0]
            dy = nearest_food[1] - head[1]
            
            if dx > 0:
                possible_directions.append(RIGHT)
            elif dx < 0:
                possible_directions.append(LEFT)
                
            if dy > 0:
                possible_directions.append(DOWN)
            elif dy < 0:
                possible_directions.append(UP)
                
            # 随机排序以避免固定模式
            random.shuffle(possible_directions)
            
            # 尝试可能的移动方向
            for direction in possible_directions:
                if self.is_safe_move(head, direction):
                    self.snake.direction = direction
                    return
                    
        # 如果没有找到食物或无法安全移动，随机选择方向
        self.random_safe_move(head)
        
    def is_safe_move(self, head, direction):
        # 计算新位置
        x, y = direction
        new_x = (head[0] + x) % GRID_WIDTH
        new_y = (head[1] + y) % GRID_HEIGHT
        new_pos = (new_x, new_y)
        
        # 检查是否撞到自己
        if new_pos in self.snake.positions[1:]:
            return False
            
        # 检查是否撞到玩家蛇
        if new_pos in self.player_snake.positions:
            return False
            
        # 避免立即反向移动
        if (self.snake.direction[0] + direction[0] == 0 and 
            self.snake.direction[1] + direction[1] == 0):
            return False
            
        return True
        
    def random_safe_move(self, head):
        # 尝试所有可能的方向
        directions = [UP, DOWN, LEFT, RIGHT]
        random.shuffle(directions)
        
        for direction in directions:
            if self.is_safe_move(head, direction):
                self.snake.direction = direction
                return
                
        # 如果没有安全的方向，保持当前方向
        return

def draw_grid(surface):
    for y in range(0, HEIGHT, GRID_SIZE):
        for x in range(0, WIDTH, GRID_SIZE):
            rect = pygame.Rect(x, y, GRID_SIZE, GRID_SIZE)
            pygame.draw.rect(surface, GRID_COLOR, rect, 1)

def draw_score(surface, player_score, enemy_score):
    # 玩家分数
    player_text = font_medium.render(f"玩家: {player_score}", True, SNAKE_HEAD_COLOR)
    player_rect = player_text.get_rect(topleft=(20, 20))
    
    # 敌人分数
    enemy_text = font_medium.render(f"敌人: {enemy_score}", True, ENEMY_HEAD_COLOR)
    enemy_rect = enemy_text.get_rect(topright=(WIDTH - 20, 20))
    
    # 绘制半透明背景
    bg_rect = pygame.Rect(10, 10, 
                         max(player_rect.width, enemy_rect.width) + 20, 
                         player_rect.height + enemy_rect.height + 20)
    s = pygame.Surface((bg_rect.width, bg_rect.height), pygame.SRCALPHA)
    s.fill(UI_BG)
    pygame.draw.rect(s, UI_BORDER, (0, 0, bg_rect.width, bg_rect.height), 2, border_radius=10)
    surface.blit(s, bg_rect)
    
    surface.blit(player_text, player_rect)
    surface.blit(enemy_text, (10, player_rect.bottom + 10))

def draw_game_over(surface, player_score, enemy_score):
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 180))
    surface.blit(overlay, (0, 0))
    
    game_over_text = font_large.render("游戏结束!", True, (255, 100, 100))
    player_text = font_medium.render(f"玩家分数: {player_score}", True, SNAKE_HEAD_COLOR)
    enemy_text = font_medium.render(f"敌人分数: {enemy_score}", True, ENEMY_HEAD_COLOR)
    restart_text = font_small.render("按 R 键重新开始", True, TEXT_COLOR)
    
    surface.blit(game_over_text, (WIDTH // 2 - game_over_text.get_width() // 2, HEIGHT // 2 - 100))
    surface.blit(player_text, (WIDTH // 2 - player_text.get_width() // 2, HEIGHT // 2 - 30))
    surface.blit(enemy_text, (WIDTH // 2 - enemy_text.get_width() // 2, HEIGHT // 2 + 20))
    surface.blit(restart_text, (WIDTH // 2 - restart_text.get_width() // 2, HEIGHT // 2 + 80))

def draw_title(surface):
    title_text = font_large.render("贪吃蛇游戏 - 带AI敌人", True, (0, 200, 200))
    surface.blit(title_text, (WIDTH // 2 - title_text.get_width() // 2, 20))

def draw_instructions(surface):
    instructions = [
        "使用方向键控制绿色蛇的移动",
        "吃到黄色食物增加分数",
        "按空格键加速",
        "按 R 键重新开始"
    ]
    
    y_offset = HEIGHT - 160
    for i, instruction in enumerate(instructions):
        text = font_small.render(instruction, True, (180, 220, 255))
        surface.blit(text, (WIDTH // 2 - text.get_width() // 2, y_offset + i * 30))

def draw_pause(surface):
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 150))
    surface.blit(overlay, (0, 0))
    
    pause_text = font_large.render("游戏暂停", True, (100, 255, 255))
    continue_text = font_medium.render("按空格键继续", True, TEXT_COLOR)
    
    surface.blit(pause_text, (WIDTH // 2 - pause_text.get_width() // 2, HEIGHT // 2 - 50))
    surface.blit(continue_text, (WIDTH // 2 - continue_text.get_width() // 2, HEIGHT // 2 + 20))

def draw_hint(surface, message, y_pos):
    hint_text = font_small.render(message, True, (255, 255, 100))
    bg_rect = pygame.Rect(WIDTH // 2 - hint_text.get_width() // 2 - 10, 
                         y_pos - 5, 
                         hint_text.get_width() + 20, 
                         hint_text.get_height() + 10)
    
    s = pygame.Surface((bg_rect.width, bg_rect.height), pygame.SRCALPHA)
    s.fill((0, 0, 0, 150))
    pygame.draw.rect(s, (255, 200, 0), (0, 0, bg_rect.width, bg_rect.height), 2, border_radius=5)
    surface.blit(s, bg_rect)
    
    surface.blit(hint_text, (WIDTH // 2 - hint_text.get_width() // 2, y_pos))

def draw_speed_boost(surface, snake):
    if snake.is_player:
        # 在蛇头周围绘制加速效果
        head = snake.get_head_position()
        center_x = head[0] * GRID_SIZE + GRID_SIZE // 2
        center_y = head[1] * GRID_SIZE + GRID_SIZE // 2
        radius = GRID_SIZE * 1.2
        
        # 创建半透明圆形
        s = pygame.Surface((radius * 2, radius * 2), pygame.SRCALPHA)
        pygame.draw.circle(s, SPEED_BOOST_COLOR, (radius, radius), radius)
        surface.blit(s, (center_x - radius, center_y - radius))

def main():
    player_snake = Snake(is_player=True)
    enemy_snake = Snake(is_player=False)
    food = Food()
    enemy_ai = EnemyAI(enemy_snake, player_snake, food)
    paused = False
    game_start_time = time.time()
    last_hint_time = time.time()
    hint_message = "按空格键加速！"
    show_hint = True
    
    while True:
        current_time = time.time()
        
        # 5秒后隐藏提示
        if show_hint and current_time - last_hint_time > 5:
            show_hint = False
            
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    if not paused:
                        # 按住空格时加速
                        player_snake.boost_speed()
                    else:
                        # 如果游戏暂停，按空格继续
                        paused = False
                elif event.key == pygame.K_r:
                    player_snake.reset()
                    enemy_snake.reset()
                    food = Food()
                    enemy_ai = EnemyAI(enemy_snake, player_snake, food)
                    paused = False
                    game_start_time = time.time()
                    last_hint_time = time.time()
                    show_hint = True
                elif not paused and player_snake.is_alive:
                    if event.key == pygame.K_UP and player_snake.direction != DOWN:
                        player_snake.direction = UP
                    elif event.key == pygame.K_DOWN and player_snake.direction != UP:
                        player_snake.direction = DOWN
                    elif event.key == pygame.K_LEFT and player_snake.direction != RIGHT:
                        player_snake.direction = LEFT
                    elif event.key == pygame.K_RIGHT and player_snake.direction != LEFT:
                        player_snake.direction = RIGHT
            elif event.type == pygame.KEYUP:
                if event.key == pygame.K_SPACE:
                    # 释放空格键时恢复原始速度
                    player_snake.reset_speed()
        
        # 填充背景
        screen.fill(BACKGROUND)
        
        # 绘制网格
        draw_grid(screen)
        
        if not paused:
            # 更新AI敌人
            enemy_ai.update()
            
            # 更新玩家蛇
            player_snake.update(food.positions, enemy_snake)
            
            # 更新敌人蛇
            enemy_snake.update(food.positions, player_snake)
        
            # 检查是否吃到食物
            head = player_snake.get_head_position()
            if head in food.positions:
                food.remove_food(head)
                player_snake.grow_to += 1
                player_snake.score += 10
                
            head = enemy_snake.get_head_position()
            if head in food.positions:
                food.remove_food(head)
                enemy_snake.grow_to += 1
                enemy_snake.score += 10
        
        # 绘制食物
        food.draw(screen)
        
        # 绘制蛇
        player_snake.draw(screen)
        enemy_snake.draw(screen)
        
        # 如果正在加速，绘制加速效果
        if pygame.key.get_pressed()[pygame.K_SPACE] and player_snake.is_alive and not paused:
            draw_speed_boost(screen, player_snake)
        
        # 绘制标题
        draw_title(screen)
        
        # 绘制分数
        draw_score(screen, player_snake.score, enemy_snake.score)
        
        # 绘制游戏说明
        draw_instructions(screen)
        
        # 显示提示
        if show_hint:
            draw_hint(screen, hint_message, HEIGHT // 4)
        
        # 游戏结束画面
        if not player_snake.is_alive:
            draw_game_over(screen, player_snake.score, enemy_snake.score)
        
        # 游戏暂停画面
        if paused and player_snake.is_alive:
            draw_pause(screen)
        
        pygame.display.flip()
        clock.tick(FPS * 3)  # 更高的帧率使动画更流畅

if __name__ == "__main__":
    main()
